;***************************************************************************************************************
; stage 2 loader for Railgun64 (0.4a +)
;
; loads bootsector to 0x9000 (first 512 bytes) to get the BPB
; loads Railgun64 to 0x10000 - 0x50000 (max size is 256k)
; then sets up higher-half paging (Kernel at 4GB - 4.5GB. 4.5GB - 6GB mapped to 2.5GB - 4GB. 0 - 3GB identity mapped)
; then sets up long mode (if available, error out otherwise)
; then transfers execution to Railgun64

; NOTE: this only works on partitions up to 8GB in size.

; data area sector start is stored at 0x7e00 (dword)
; partition sector start is stored at 0x7e10 (dword)


[BITS 16]
[ORG 0x9000]

START:
	; jump over BPB
	jmp	BPB_END

;***************************************************************************************************************
; the Bios Parameter Block
; make sure it is exactly 3 bytes in
times	3 - ($ - $$) db 0

BPB_OEMName			db "mkdosfs",0
BPB_BytesPerSector		dw 0x0200
BPB_SectorsPerCluster		db 0x08
BPB_ReservedSectors		dw 0x0020
BPB_NumberOfFATs		db 0x02
BPB_RootEntryCount		dw 0	; 0 for FAT32
BPB_TotalSectors16		dw 0	; 0 for FAT32
BPB_MediaDescriptor		db 0xf8	; 0xf8 for fixed disks
BPB_FATSize16			dw 0	; 0 for FAT32
BPB_SectorsPerTrack		dw 0x003f
BPB_NumberOfHeads		dw 0x00ff
BPB_HiddenSectors		dd 0
BPB_TotalSectors32		dd 0x001fddc2
BPB_FATSize32			dd 0x000007f4
BPB_ExtFlags			dw 0
BPB_VersionNumber		dw 0
BPB_RootCluster			dd 0x00000002
BPB_FSInfo			dw 0x0001
BPB_BootSectorBackUp		dw 0x0006
BPB_Reserved0			dd 0
BPB_Reserved1			dd 0
BPB_Reserved2			dd 0
BPB_DriveNumber			db 0
BS_Reserved3			db 0
BS_BootSig			db 0
BS_VolumeID			dd 0
BS_VolumeLabel			db "Test       "
BS_FileSystemType		db "FAT32   "

times 512 - ($ - $$) db 0	; take us to the end of the sector

;***************************************************************************************************************

BPB_END:
	; save drive number
	mov	[BPB_DriveNumber], dl
	push	dx

	; load first sector of partition to 0x9000
	; this is to fill in the above values ^
	mov	eax, 0
	mov	cx, 1
	mov	si, 0x9000
	mov	bx, 0x0000
	call	ReadDrive
	pop	dx
	mov	[BPB_DriveNumber], dl

LOAD_ROOT_DIR_CLUSTER:
	; load the cluster containing the root dir
	mov	eax, [BPB_RootCluster]
	mov	bx, 0x0000
	mov	si, 0x8000
	call	LoadCluster

FIND_RAILGUN_DIR:
	; need to find the railgun directory "Railgun"
	; setup for loop through directory entries
	mov	si, 0x8000	; point si to start of entries
	mov	cx, 128		; max 128 directories

.loop:
	; the main search loop
	; si = start of this dir entry
	; compare the directory name
	push	cx		; save loop counter
	mov	di, msgDirName	; get the directory name
	mov	cx, 11		; 11 chars per short file name
	call	CompString	; compare them
	jnc	.found_dir	; if carry not set, strings are equal
	pop	cx		; restore loop counter

	; not this entry, move onto next
	add	si, 32
	loop	.loop

	; loop fell through
	jmp	NO_RAILGUN_DIR

.found_dir:
	; found the directory, si points to it
	add	sp, 2		; restore stack
	mov	ax, [si + 20]	; get high part of cluster number
	shl	eax, 16		; shift eax left by 16 bits
	mov	ax, [si + 26]	; get low part of cluster number

	; eax now contains cluster number of Railgun directory
	; load this cluster to 0x8000 and find the Railgun64.sys file
	mov	bx, 0x0000
	mov	si, 0x8000
	call	LoadCluster

FIND_RAILGUN_FILE:
	; need to find the railgun file "Railgun64.sys"
	; setup for loop through directory entries
	mov	si, 0x8000	; point si to start of entries
	mov	cx, 128		; max 128 directories

.loop:
	; the main search loop
	; si = start of this dir entry
	; compare the directory name
	push	cx		; save loop counter
	mov	di, msgFileName	; get the directory name
	mov	cx, 11		; 11 chars per short file name
	call	CompString	; compare them
	jnc	.found_dir	; if carry not set, strings are equal
	pop	cx		; restore loop counter

	; not this entry, move onto next
	add	si, 32
	loop	.loop

	; loop fell through
	jmp	NO_RAILGUN64

.found_dir:
	; found the directory, si points to it
	add	sp, 2		; restore stack
	mov	ax, [si + 20]	; get high part of cluster number
	shl	eax, 16		; shift eax left by 16 bits
	mov	ax, [si + 26]	; get low part of cluster number
	push	eax		; save it for later

LOAD_RAILGUN:
	; eax contains starting cluster number of file
	; load it to 0x1000:0x0000
	mov	bx, 0x1000
	mov	si, 0x0000
	call	LoadCluster

	; setup for cluster chain load loop
	pop	eax		; restore starting cluster
	mov	bx, 0x1000	; set starting segment
	mov	si, 0x1000	; set starting offset (after first cluster)

.loop:
	; the main cluster load loop
	; get next cluster number
	; find which sector contains the part of the FAT containing this cluster entry
	; eax = previous cluster number
	; bx:si = pointer to next load location
	mov	ecx, 128			; number of FAT entries per sector
	xor	edx, edx			; clear edx
	div	ecx				; get sector offset in eax, entry index in edx
	shl	edx, 2				; multiply edx by 4 to get byte offset
	movzx	ecx, WORD [BPB_ReservedSectors]	; get start sector of FAT
	add	eax, ecx			; add to offset to get the sector to load

	; now load the FAT sector to 0x0600
	push	ebx				; save file target segment
	push	esi				; save file target offset
	push	edx				; save index
	mov	bx, 0x0000			; FAT sector target segment
	mov	si, 0x0600			; FAT sector target offset
	mov	ecx, 1				; want 1 sector
	call	ReadDrive			; read in the sector

	; sector loaded, now get the next cluster number from it
	pop	esi				; get the entry offset
	add	si, 0x0600			; add the base address of the sector
	mov	eax, [si]			; get the next cluster number

	; check cluster number for end of file
	; cluster number is in eax
	and	eax, 0x0fffffff			; mask top 4 bits
	cmp	eax, 0x0fffffef			; check for EOF signature
	jg	RAILGUN_LOADED			; leave if present

	; cluster chain not finished
	; load cluster
	; eax = cluster number to load
	pop	esi				; restore file cluster pointer
	pop	ebx				; restore file cluster segment
	push	ebx				; save segment pointer again
	push	esi				; save offset pointer again
	push	eax				; save cluster number
	call	LoadCluster			; load the cluster
	pop	eax				; restore cluster number
	pop	esi				; restore offset pointer
	pop	ebx				; restore segment pointer

	; increase destination pointer
	add	si, 0x1000			; increase si by 0x1000
	or	si, si				; check if it's zeroed (overflowed)
	jnz	.continue
	add	bx, 0x1000			; if it did zero, add 0x1000 to segment

.continue:
	; eax = last segment loaded
	; si = destination offset for next cluster
	; bx = destination segment for next cluster
	; now load the next one
	jmp	.loop

RAILGUN_LOADED:
	; kernel loaded, proceed to setup system for use
	; open A20 gate
	; use the fast A20 open option
	mov	ax, 0x2401	; ax = 0x2401 for fast A20
	int	0x15		; call it
	jc	NO_A20		; error out on error

	; create a memory map at 0x0000:0xe000
	; and memory map entry count at 0x0000:0xfff0
	mov	ax, 0x0000
	mov	es, ax
	mov	di, 0xe000
	call	do_e820

	; save entry count
	mov	[0xfff0], bp

	; check for long mode availability
	; use CPUID to check if long mode is supported
	; check for cpuid extentions
	mov 	eax, 0x80000000 	; Set the A-register to 0x80000000.
	cpuid                  		; CPU identification.
    	cmp 	eax, 0x80000001		; Compare the A-register with 0x80000001.
    	jb 	NO_LONG_MODE		; It is less, there is no long mode.

	; check for long mode
	mov	eax, 0x80000001		; Set the A-register to 0x80000001.
	cpuid                  		; CPU identification.
	test	edx, 1 << 29      	; Test if the LM-bit, which is bit 29, is set in the D-register.
	jz	NO_LONG_MODE         	; They aren't, there is no long mode.

	; long mode is supported, continue
	; setup for mininal protected mode environment
	cli				; disable interrupts
	lgdt	[prmd_gdt_desc]		; load basic GDT
	mov	eax, cr0		; get cr0
	or	al, 1			; set protected mode switch
	mov	cr0, eax		; save it back in cr0
	jmp	0x08:START32		; jump to 32 bit segment to enable 32 bit protected mode

;*************************************************************************************************
[BITS 32]
START32:
	; setup the data segment
	mov	ax, 0x10
	mov	ds, ax

	; now enable long mode
	; need to create basic page tables (identity map the first 2 mb)
	; put the page table at 0x1000
	mov	edi, 0x1000		; put the plm4 address in edi
	mov	cr3, edi		; give cr3 the location of the page table pml4

	; create a pml4 entry
	mov	eax, 0x1b + 0x2000	; flags + address of pdpe entry
	xor	ebx, ebx		; set high half to 0
	mov	[edi], eax		; write low half
	mov	[edi + 4], ebx		; write high half

	; create a pdpe entry at 0x2000
	mov	edi, 0x2000
	mov	eax, 0x301b		; flags + address
	mov	[edi], eax		; write low half
	mov	[edi + 4], ebx		; write high half

	; create 2 pde entries at 0x3000
	mov	edi, 0x3000
	mov	eax, 0x9b		; flags + 0 = identity map 1st 2mb
	mov	[edi], eax		; write low half
	mov	[edi + 4], ebx		; write high half

	mov	eax, 0x20009b		; flags + 0x200000 = identity map 2nd 2mb
	mov	[edi + 8], eax		; write low half
	mov	[edi + 12], ebx		; write high half

	; enable PAE
	mov	eax, cr4
	or	al, 1 << 5
	mov	cr4, eax

	; enable long mode
	mov	ecx, 0xc0000080
	rdmsr
	or	ax, 1 << 8
	wrmsr

	; enable paging
	mov	eax, cr0
	or	eax, 1 << 31
	mov	cr0, eax

	; load the long mode GDT
	lgdt	[long_gdt_desc]

	; jump to long mode code
	jmp	0x08:START64

;*************************************************************************************************
[BITS 64]
START64:
	; now in long mode
	; setup temporary segments
	mov	ax, 0x10
	mov	ds, ax
	mov	es, ax
	mov	ax, 0
	mov	ss, ax

	; now create proper page table at 0x100000 (1MB)
	; identity map 0 - 3GB (temporary) (write-through cached)
	; identity map map 3GB - 4GB (strong uncached, global)
	; map 4GB - 4.5GB to 0 - 0.5GB (write-through cached, global)

CREATE_PAGE_TABLE:
	; create a pml4 entry
	mov	rdi, 0x100000
	mov	rax, 0x101000 + 7	; flags + address of first 512 pdpe entries
	mov	[rdi], rax

	; setup for creating 512 pdpe entries at 0x101000 - 0x101fff
	mov	rdi, 0x101000		; base address of pdpe entries
	mov	rcx, 512		; want 512 entries
	mov	rax, 0x102000 + 7	; flags + base address of pdp entries (0x102000)

.pdpe_loop:
	; put the entry in, etc
	mov	[rdi], rax		; put rax at rdi
	add	rax, 0x1000		; add 4096 to eax to point it to next pdp block
	add	rdi, 8			; increment edi by 8 to point to next pdpe entry
	loop	.pdpe_loop		; do this 512 times

	; done pdpe
	; now do first 1536 pdp entries (identity mapped 0 - 3GB)
	; this is temporary - allows trampoline code to work for APs until all are initialised
	mov	rdi, 0x102000		; rdi points to start of pdp entries
	mov	rax, 0x8b		; base entry points to 2mb page starting at address 0
	mov	rcx, 1536		; want 1536 entries

.pdp_3GB_loop:
	; put the entry in, etc
	mov	[rdi], rax		; put the entry at rdi
	;mov	rax, 0			; clear all but first page
	add	rax, 0x200000
	add	rdi, 8			; point rdi to next entry in table
	loop	.pdp_3GB_loop		; do this 1024 times

	; done pdp 3GB identity map
	; now do kernel 512MB
	; mapped to 0 - 512MB
	mov	rdi, 0x106000		; rdi points to next block of pdp entries
	mov	rax, 0x18b		; base entry points to 2mb page starting at address 0 with flags
	mov	rcx, 256		; want 256 of these

.pdp_k_loop:
	; put the entry in, etc
	mov	[rdi], rax		; put the entry at rdi
	add	rax, 0x200000		; increase rax by 2mb to point to next page
	add	rdi, 8			; point rdi to next entry in table
	loop	.pdp_k_loop		; do this 256 times

	; done pdp kernel map
	; now do pdp memory mapped IO map
	; identity map 3GB - 4GB
	mov	rdi, 0x105000		; set rdi to point to block of pdp entries
	mov	rax, 0xc000019b		; set rax to 3GB with flags
	mov	rcx, 512		; want 512 of these

.pdp_m_loop:
	; put the entry in, etc
	mov	[rdi], rax		; put the entry at rdi
	add	rax, 0x200000		; increase rax by 2mb to point to next page
	add	rdi, 8			; point rdi to next entry in table
	loop	.pdp_m_loop		; do this 768 times

	; done, set cr3 to point to the map
	mov	rax, 0x100000
	mov	cr3, rax

ENABLE_FEATURES:
	; enable SSE and the like
	mov 	rax, cr0
	and 	ax, 0xfffb		; clear coprocessor emulation CR0.EM
	or 		ax, 0x2			; set coprocessor monitoring CR0.MP
	mov 	cr0, rax
	mov 	rax, cr4
	or 		ax, 3 << 9		; set CR4.OSFXSR and CR4.OSXMMEXCPT at the same time
	mov 	cr4, rax

	; enable paging global pages
	mov	rax, cr4
	or	rax, (1 << 7)
	mov	cr4, rax


GOTO_CODE:
	mov	rax, 0x100011000
	jmp	rax


;**************************************************************************************************************
; 16 bit error outs
[BITS 16]

NO_A20:
	mov	si, msgNoA20
	call	PrintString
	cli
	hlt

NO_RAILGUN64:
	mov	si, msgNoRailgun
	call	PrintString
	cli
	hlt

NO_RAILGUN_DIR:
	mov	si, msgNoRailgunDir
	call	PrintString
	cli
	hlt

NO_LONG_MODE:
	mov	si, msgNoLongMode
	call	PrintString
	cli
	hlt

;***************************************************************************************************************
; functions (16 bit)

;*******************************
; compare 2 strings
; si = ptr to string 1
; di = ptr to string 2
; cx = length of string
; sets carry on fail

CompString:
	; si and di point to the chars
	; save di & si
	push	di
	push	si

.loop:
	; check the char
	mov		al, [si]
	mov		bl, [di]
	inc		si
	inc		di
	cmp		al, bl
	jne		.fail
	loop	.loop

.success:
	pop		si
	pop		di
	clc
	ret

.fail:
	pop		si
	pop		di
	stc
	ret

;*******************************
; print a 32 bit number (unsigned)
; eax = number
; ecx = base (2 - 16)

PrintNumber:
	; setup a null-terminated array on the stack
	; setup so that default array is "0",0 (char for 0, with nullterminator)
	; to account for values of 0 being printed
	; point si to a word before the end
	push	WORD 0x0030
	mov	si, sp
	sub	sp, 12

	; check value is not zero
	or	eax, eax
	jz	.loop

	; number is not zero, increase si to remove "0" char
	inc	si

.loop:
	; get the next character's value
	xor	edx, edx
	div	ecx

	; check if finished (quotient = 0, remainder = 0)
	or	eax, eax	; quotient
	jnz	.not_finished
	or	edx, edx	; remainder
	jnz	.not_finished
	jmp	.print

.not_finished:
	; decrease si to point to previous byte - create characters in reverse
	dec	si

	; check if value is larger than 10
	cmp	dx, 10
	jl	.smaller_10

	; value is larger than 10, so add offset to capital letters from
	; numbers - 10 (ASCII code)
	add	dx, 0x07

.smaller_10:
	; add base value for ASCII numbers 0 - 9
	add	dx, 0x30

	; save the value
	mov	[si], dl

	; loop
	jmp	.loop

.print:
	; print the generated string
	call	PrintString

.finish:
	; restore stack + return
	add	sp, 14
	ret

;*******************************
; print a string
; si = pointer to string to print

PrintString:
	; setup registers for video out mode
	mov	ah, 0x0e
	mov	bh, 0x00
	mov	bl, 0x07

.loop:
	; get char
	lodsb		; same as "mov	al, [si] + inc	si"

	; check for nullchar
	or	al, al
	jz	.finish

	; if not null, print it & repeat
	int	0x10
	jmp	.loop

.finish:
	ret

;*******************************
; load a cluster
; eax = cluster to load
; bx:si = target start address

LoadCluster:
	; calculate start LBA
	; first sector of cluster = ((N - 2) * sectors_per_cluster) + first data sector
	sub	eax, 2					; N - 2
	movzx	ecx, BYTE [BPB_SectorsPerCluster]	; load sectors per cluster into cl
	mul	ecx					; multiply
	movzx	ecx, BYTE [BPB_SectorsPerCluster]	; get sectors per cluster again (needed?)
	add	eax, [0x7e00]				; where the first data sector value is stored

	; call the drive read function
	call	ReadDrive

.finish:
	ret

;*******************************
; new drive read function
; eax = start lba
; cx = sector count
; bx:si = target start address

ReadDrive:
	; fetch the drive number
	mov	dl, [BPB_DriveNumber]

	; add sector offset for partition start
	add	eax, [0x7e10]

	; put DAP values onto the stack
	; must be done in REVERSE order
	push	DWORD 0		; DAP sector start high 32 bits
	push	eax		; DAP sector start low 32 bits
	push	bx		; DAP segment
	push	si		; DAP offset
	push	cx		; DAP sector count
	push	WORD 0x0010	; DAP size + unused

	; save the pointer of the top of the stack
	push	sp

	; setup for the loop
	mov	di, 10		; 10 attempts

.loop:
	; decrease loop counter
	dec	di

	; check if counter = 0
	or	di, di
	jz	.fail

	; attempt to reset the drive
	; setup for drive reset
	mov	ah, 0

	; call the interrupt
	int	0x13

	; if reset failed, try again
	jc	.loop

	; drive reset was success, now try the read
	; setup for extended read
	pop	si		; get the pointer to the DAP
	push	ds		; save ds
	mov	ax, ss		; get the segment of the DAP
	mov	ds, ax
	mov	ah, 0x42	; extended read

	; call the interrrupt
	int	0x13

	; restore ds
	pop	ds

	; if read failed, try again
	jc	.loop

	; restore the stack
	add	sp, 16

.finish:
	; finished - return
	ret

.fail:
	; failed - print error message and halt
	mov	si, msgReadError
	call	PrintString
	cli
	hlt

;*************************************************************************************************
; function to get the memory map via BIOS int 0x15:e820
; used from OSDev wiki with thanks :) http://wiki.osdev.org/Detecting_Memory_%28x86%29#Getting_an_E820_Memory_Map
; use the INT 0x15, eax= 0xE820 BIOS function to get a memory map
; inputs: es:di -> destination buffer for 24 byte entries
; outputs: bp = entry count, trashes all registers except esi

do_e820:
	xor 	ebx, ebx		; ebx must be 0 to start
	xor	bp, bp			; keep an entry count in bp
	mov 	edx, 0x0534D4150	; Place "SMAP" into edx
	mov 	eax, 0xe820
	mov 	[di + 20], dword 1	; force a valid ACPI 3.X entry
	mov 	ecx, 24			; ask for 24 bytes
	int 	0x15
	jc 	short .failed		; carry set on first call means "unsupported function"
	mov 	edx, 0x0534D4150	; Some BIOSes apparently trash this register?
	cmp 	eax, edx		; on success, eax must have been reset to "SMAP"
	jne 	short .failed
	test 	ebx, ebx		; ebx = 0 implies list is only 1 entry long (worthless)
	je 	short .failed
	jmp 	short .jmpin

.e820lp:
	mov 	eax, 0xe820		; eax, ecx get trashed on every int 0x15 call
	mov 	[di + 20], dword 1	; force a valid ACPI 3.X entry
	mov 	ecx, 24			; ask for 24 bytes again
	int 	0x15
	jc 	short .e820f		; carry set means "end of list already reached"
	mov 	edx, 0x0534D4150	; repair potentially trashed register

.jmpin:
	jcxz 	.skipent		; skip any 0 length entries
	cmp 	cl, 20			; got a 24 byte ACPI 3.X response?
	jbe 	short .notext
	test 	byte [di + 20], 1	; if so: is the "ignore this data" bit clear?
	je 	short .skipent

.notext:
	mov 	ecx, [di + 8]		; get lower dword of memory region length
	or 	ecx, [di + 12]		; "or" it with upper dword to test for zero
	jz 	.skipent		; if length qword is 0, skip entry
	inc 	bp			; got a good entry: ++count, move to next storage spot
	add 	di, 24

.skipent:
	test	ebx, ebx		; if ebx resets to 0, list is complete
	jne 	short .e820lp

.e820f:
	ret

.failed:
	; error occured
	mov	si, msgNoMemoryMap
	call	PrintString
	cli
	hlt

;****************************************************************************************
; data

msgDirName		db "RAILGUNX   ",0 
msgFileName		db "RAILGU~1SYS",0
msgNoRailgunDir		db 13,10,"ERROR: could not find /RailgunX directory",0
msgNoRailgun		db 13,10,"ERROR: could not find /Railgun/Railgun64.sys",0
msgReadError		db 13,10,"ERROR: could not read from drive",0
msgNoA20		db 13,10,"ERROR: could not open A20 gate",0
msgNoMemoryMap		db 13,10,"ERROR: could not generate memory map",0
msgNoLongMode		db 13,10,"ERROR: CPU does not support 64 bit mode",0

;****************************************************************************************
; protected mode flat GDT

prmd_gdt_table:
	; null
	dq 0

	; kernel code
	dw 0xFFFF
	dw 0x0000
	db 0x00
	db 10011010b
	db 11001111b
	db 0x00

	; kernel data
	dw 0xFFFF
	dw 0x0000
	db 0x00
	db 10010010b
	db 11001111b
	db 0x00

prmd_gdt_desc:
	dw $ - prmd_gdt_table - 1
	dd prmd_gdt_table

;****************************************************************************************
; long mode flat GDT

long_gdt_table:
	; null
	dq 0

	; kernel code
	dw 0
	dw 0
	db 0
	db 10011000b
	db 00100000b
	db 0

	; kernel data
	dw 0
	dw 0
	db 0
	db 10010000b
	db 00000000b
	db 0

long_gdt_desc:
	dw $ - long_gdt_table - 1
	dq long_gdt_table


;****************************************************************************************
times 0x1000 - ($ - $$) db 0






